from core.function import Function
from core.cuda import cuda
from implement.functions.tensor_opera.get_item_grad import GetItemGrad


# 切片
class GetItem(Function):
    def __init__(self, slices):
        """
        初始化 GetItem 函数。

        Parameters:
            slices: 切片信息。

        Returns:
            None
        """
        self.slices = slices

    def forward(self, x):
        """
        前向传播函数，实现对输入数据的切片操作。

        Parameters:
            x: 输入数据。

        Returns:
            y: 切片后的数据。
        """
        x = cuda.to_array(x)
        y = x[self.slices]
        return y

    def backward(self, gy):
        """
        反向传播函数，计算梯度。

        Parameters:
            gy: 相对于前向传播结果的梯度。

        Returns:
            gx: 相对于输入数据的梯度。
        """
        x, = self.inputs
        f = GetItemGrad(self.slices, x.shape)
        return f(gy)

    def get_item(self, x):
        """
        获取切片结果的快捷方法。

        Parameters:
            x: 输入数据。

        Returns:
            y: 切片结果。
        """
        return self(x)


if __name__ == '__main__':
    from core.variable import Variable
    # 测试
    x = Variable([[1, 2, 3], [4, 5, 6]])
    y = x[1]
    print(y.data)

    y = x[:, 2]
    print(y.data)
